import os
import sys

from fastapi import FastAPI, UploadFile, File
import click as click
import uvicorn as uvicorn

from settings import get_settings

sets = get_settings()
sys.path.append(sets.OCRPATH)
sys.path.append(sets.OCRPATH_F)

import paddle

from ppocr.data import create_operators, transform
from ppocr.modeling.architectures import build_model
from ppocr.postprocess import build_post_process
from ppocr.utils.save_load import init_model
import tools.program as program
import numpy as np

os.environ["FLAGS_allocator_strategy"] = 'auto_growth'

app = FastAPI(title="ocr project", openapi_url=f"/api/v1/openapi.json")

SUCCESS = "SUCCESS"
FAILED = "FAILED"


class ReturnInfo(object):
    def __init__(self):
        self.status = FAILED  # SUCCESS FAILED
        self.msg_code = ''
        self.msg = ''
        self.data = {}

    def todict(self):
        # __dict__ 是字典，不进行转化无法输出
        return self.__dict__


@app.on_event("startup")
async def startup():
    """
    在所有 startup 事件处理程序完成 之前，您的应用程序将不会开始接收请求 。
    :return:
    """
    config, device, logger, vdl_writer = program.preprocess(yum_path=sets.YAMLPATH)
    global_config = config['Global']

    # build post process
    post_process_class = build_post_process(config['PostProcess'],
                                            global_config)

    # build model
    if hasattr(post_process_class, 'character'):
        config['Architecture']["Head"]['out_channels'] = len(getattr(post_process_class, 'character'))

    model = build_model(config['Architecture'])

    init_model(config, model, logger)

    # create data ops
    transforms = []
    for op in config['Eval']['dataset']['transforms']:
        op_name = list(op)[0]
        if 'Label' in op_name:
            continue
        elif op_name in ['RecResizeImg']:
            op[op_name]['infer_mode'] = True
        elif op_name == 'KeepKeys':
            if config['Architecture']['algorithm'] == "SRN":
                op[op_name]['keep_keys'] = [
                    'image', 'encoder_word_pos', 'gsrm_word_pos',
                    'gsrm_slf_attn_bias1', 'gsrm_slf_attn_bias2'
                ]
            else:
                op[op_name]['keep_keys'] = ['image']
        transforms.append(op)
    global_config['infer_mode'] = True
    ops = create_operators(transforms, global_config)
    model.eval()
    app.state.ops = ops
    app.state.model = model
    app.state.post_process_class = post_process_class


@app.on_event("shutdown")
async def shutdown():
    """
    要添加应在应用程序关闭时运行的功能，请使用事件声明它 "shutdown"
    :return:
    """
    pass


@app.post("/ocr/verification_code")
async def down_home(file: UploadFile = File(...)):
    return_info = ReturnInfo()
    ops = app.state.ops
    model = app.state.model
    post_process_class = app.state.post_process_class
    img = await file.read()
    data = {'image': img}
    batch = transform(data, ops)
    images = np.expand_dims(batch[0], axis=0)
    images = paddle.to_tensor(images)
    preds = model(images)
    post_result = post_process_class(preds)
    return_info.status = SUCCESS
    return_info.data = {"code": post_result[0][0], "Precision": str(post_result[0][1])}
    return return_info.todict()


@click.command()
@click.option('--port', default=8001, help='set port')
def start(port):
    """
    :param port: 传入端口号，默认为8001
    :param tags: main为调试正确的服务，test为调试服务
    :param app_name:
    :return:
    """
    uvicorn.run(app='main:app', host='0.0.0.0', port=port, reload=True)


# python main.py --port 8001
if __name__ == '__main__':
    start()
